﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;
using System.Runtime.InteropServices;

namespace SpectralCamera.Extend
{
    public class MyUdp : IDisposable
    {
        private readonly UdpClient _udpClient;
        private readonly IPEndPoint _sendEndPoint;
        private readonly bool _isRemote;
        private bool _isCallbackReady;
        private bool _isSyncRecv;
        private byte[] _syncRecvBuf = new byte[1024 * 1024];
        private int _syncRecvLen;
        private CancellationTokenSource _cts = new CancellationTokenSource();

        public delegate void RecvCallback(MessageUDP msg);
        public event RecvCallback OnReceived;

        public MyUdp(short recvPort, string sendIp, short sendPort, bool isRmt)
        {
            _isRemote = isRmt;
            _sendEndPoint = new IPEndPoint(IPAddress.Parse(sendIp), sendPort);
            _udpClient = new UdpClient(new IPEndPoint(IPAddress.Any, recvPort));
            _udpClient.Client.ReceiveBufferSize = 100 * 1024 * 1024;
            _ = Task.Run(() => ReceiveLoop(_cts.Token));
        }

        public void RegisterRecvCallback(RecvCallback cb)
        {
            OnReceived += cb;
            _isCallbackReady = true;
        }

        public async Task<int> RequestExAsync(byte[] sendBuf, int sendLen, byte[] recvBuf, int delayTime, int recvTimeout)
        {
            _isSyncRecv = true;
            Array.Clear(_syncRecvBuf, 0, _syncRecvBuf.Length);

            await SendAsync(sendBuf, sendLen);
            await Task.Delay(delayTime);

            Array.Copy(_syncRecvBuf, recvBuf, _syncRecvLen);
            _isSyncRecv = false;
            return _syncRecvLen;
        }

        public async Task<int> SendAsync(byte[] data, int length)
        {
#if PRINT_FOR_DEBUG
        Console.WriteLine($"Send: {BitConverter.ToString(data, 0, Math.Min(20, length))}");
#endif
            return await _udpClient.SendAsync(data, length, _sendEndPoint);
        }

        private async void ReceiveLoop(CancellationToken token)
        {
            while (!token.IsCancellationRequested)
            {
                try
                {
                    var result = await _udpClient.ReceiveAsync().WithCancellation(token);
                    ProcessReceive(result.Buffer, result.RemoteEndPoint);
                }
                catch (OperationCanceledException)
                {
                    break;
                }
            }
        }

        private void ProcessReceive(byte[] buffer, IPEndPoint remoteEp)
        {
            if (buffer.Length <= 12) return;

            var msg = new MessageUDP
            {
                Head = BitConverter.ToUInt32(buffer, 0),
                Rsv = buffer[4],
                ModuleType = buffer[5],
                ModuleAddr = buffer[6],
                ModuleGrp = buffer[7],
                Cmd = (ushort)((buffer[8] << 8) | buffer[9]),
                Len = (ushort)((buffer[10] << 8) | buffer[11])
            };

            if (msg.Len > 0 && msg.Len < 0xFFFF)
            {
                if (_isSyncRecv)
                {
                    _syncRecvLen = buffer.Length;
                    Buffer.BlockCopy(buffer, 0, _syncRecvBuf, 0, buffer.Length);
                }
                else if (_isCallbackReady)
                {
                    msg.Data = new byte[msg.Len];
                    Buffer.BlockCopy(buffer, 12, msg.Data, 0, msg.Len);
                    OnReceived?.Invoke(msg);
                }
            }
        }

        public void Dispose()
        {
            _cts.Cancel();
            _udpClient?.Dispose();
        }
    }

    public static class UdpExtensions
    {
        public static async Task<UdpReceiveResult> WithCancellation(
            this Task<UdpReceiveResult> task, CancellationToken token)
        {
            var tcs = new TaskCompletionSource<bool>();
            using (token.Register(() => tcs.TrySetResult(true)))
            {
                if (task != await Task.WhenAny(task, tcs.Task))
                    throw new OperationCanceledException(token);
            }
            return await task;
        }
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    public struct MessageUDP
    {
        public uint Head;
        public byte Rsv;
        public byte ModuleType;
        public byte ModuleAddr;
        public byte ModuleGrp;
        public ushort Cmd;
        public ushort Len;
        public byte[] Data;
    }
}
